/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.sdk.transforms;



import com.bff.gaia.unified.sdk.coders.*;

import com.bff.gaia.unified.sdk.transforms.Combine.AccumulatingCombineFn;

import com.bff.gaia.unified.sdk.transforms.Combine.AccumulatingCombineFn.Accumulator;

import com.bff.gaia.unified.sdk.transforms.Combine.PerKey;

import com.bff.gaia.unified.sdk.transforms.display.DisplayData;

import com.bff.gaia.unified.sdk.transforms.windowing.GlobalWindow;

import com.bff.gaia.unified.sdk.transforms.windowing.GlobalWindows;

import com.bff.gaia.unified.sdk.util.NameUtils;

import com.bff.gaia.unified.sdk.util.common.ElementByteSizeObserver;

import com.bff.gaia.unified.sdk.values.KV;

import com.bff.gaia.unified.sdk.values.PCollection;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.Lists;

import com.bff.gaia.unified.sdk.coders.Coder;

import com.bff.gaia.unified.sdk.coders.CoderException;

import com.bff.gaia.unified.sdk.coders.CoderRegistry;

import com.bff.gaia.unified.sdk.coders.CustomCoder;

import com.bff.gaia.unified.sdk.coders.ListCoder;



import javax.annotation.Nullable;

import java.io.IOException;

import java.io.InputStream;

import java.io.OutputStream;

import java.io.Serializable;

import java.util.*;



import static com.bff.gaia.unified.vendor.guava.com.google.common.base.Preconditions.checkArgument;



/**

 * {@code PTransform}s for finding the largest (or smallest) set of elements in a {@code

 * PCollection}, or the largest (or smallest) set of values associated with each key in a {@code

 * PCollection} of {@code KV}s.

 */

public class Top {



  private Top() {

    // do not instantiate

  }



  /**

   * Returns a {@code PTransform} that takes an input {@code PCollection<T>} and returns a {@code

   * PCollection<List<T>>} with a single element containing the largest {@code count} elements of

   * the input {@code PCollection<T>}, in decreasing order, sorted using the given {@code

   * Comparator<T>}. The {@code Comparator<T>} must also be {@code Serializable}.

   *

   * <p>If {@code count} {@code >} the number of elements in the input {@code PCollection}, then all

   * the elements of the input {@code PCollection} will be in the resulting {@code List}, albeit in

   * sorted order.

   *

   * <p>All the elements of the result's {@code List} must fit into the memory of a single machine.

   *

   * <p>Example of use:

   *

   * <pre>{@code

   * PCollection<Student> students = ...;

   * PCollection<List<Student>> top10Students =

   *     students.apply(Top.of(10, new CompareStudentsByAvgGrade()));

   * }</pre>

   *

   * <p>By default, the {@code Coder} of the output {@code PCollection} is a {@code ListCoder} of

   * the {@code Coder} of the elements of the input {@code PCollection}.

   *

   * <p>If the input {@code PCollection} is windowed into {@link GlobalWindows}, an empty {@code

   * List<T>} in the {@link GlobalWindow} will be output if the input {@code PCollection} is empty.

   * To use this with inputs with other windowing, either {@link Combine.Globally#withoutDefaults

   * withoutDefaults} or {@link Combine.Globally#asSingletonView asSingletonView} must be called.

   *

   * <p>See also {@link #smallest} and {@link #largest}, which sort {@code Comparable} elements

   * using their natural ordering.

   *

   * <p>See also {@link #perKey}, {@link #smallestPerKey}, and {@link #largestPerKey}, which take a

   * {@code PCollection} of {@code KV}s and return the top values associated with each key.

   */

  public static <T, ComparatorT extends Comparator<T> & Serializable>

      Combine.Globally<T, List<T>> of(int count, ComparatorT compareFn) {

    return Combine.globally(new TopCombineFn<>(count, compareFn));

  }



  /**

   * Returns a {@code PTransform} that takes an input {@code PCollection<T>} and returns a {@code

   * PCollection<List<T>>} with a single element containing the smallest {@code count} elements of

   * the input {@code PCollection<T>}, in increasing order, sorted according to their natural order.

   *

   * <p>If {@code count} {@code >} the number of elements in the input {@code PCollection}, then all

   * the elements of the input {@code PCollection} will be in the resulting {@code PCollection}'s

   * {@code List}, albeit in sorted order.

   *

   * <p>All the elements of the result {@code List} must fit into the memory of a single machine.

   *

   * <p>Example of use:

   *

   * <pre>{@code

   * PCollection<Integer> values = ...;

   * PCollection<List<Integer>> smallest10Values = values.apply(Top.smallest(10));

   * }</pre>

   *

   * <p>By default, the {@code Coder} of the output {@code PCollection} is a {@code ListCoder} of

   * the {@code Coder} of the elements of the input {@code PCollection}.

   *

   * <p>If the input {@code PCollection} is windowed into {@link GlobalWindows}, an empty {@code

   * List<T>} in the {@link GlobalWindow} will be output if the input {@code PCollection} is empty.

   * To use this with inputs with other windowing, either {@link Combine.Globally#withoutDefaults

   * withoutDefaults} or {@link Combine.Globally#asSingletonView asSingletonView} must be called.

   *

   * <p>See also {@link #largest}.

   *

   * <p>See also {@link #of}, which sorts using a user-specified {@code Comparator} function.

   *

   * <p>See also {@link #perKey}, {@link #smallestPerKey}, and {@link #largestPerKey}, which take a

   * {@code PCollection} of {@code KV}s and return the top values associated with each key.

   */

  public static <T extends Comparable<T>> Combine.Globally<T, List<T>> smallest(int count) {

    return Combine.globally(new TopCombineFn<>(count, new Reversed<T>()));

  }



  /**

   * Returns a {@code PTransform} that takes an input {@code PCollection<T>} and returns a {@code

   * PCollection<List<T>>} with a single element containing the largest {@code count} elements of

   * the input {@code PCollection<T>}, in decreasing order, sorted according to their natural order.

   *

   * <p>If {@code count} {@code >} the number of elements in the input {@code PCollection}, then all

   * the elements of the input {@code PCollection} will be in the resulting {@code PCollection}'s

   * {@code List}, albeit in sorted order.

   *

   * <p>All the elements of the result's {@code List} must fit into the memory of a single machine.

   *

   * <p>Example of use:

   *

   * <pre>{@code

   * PCollection<Integer> values = ...;

   * PCollection<List<Integer>> largest10Values = values.apply(Top.largest(10));

   * }</pre>

   *

   * <p>By default, the {@code Coder} of the output {@code PCollection} is a {@code ListCoder} of

   * the {@code Coder} of the elements of the input {@code PCollection}.

   *

   * <p>If the input {@code PCollection} is windowed into {@link GlobalWindows}, an empty {@code

   * List<T>} in the {@link GlobalWindow} will be output if the input {@code PCollection} is empty.

   * To use this with inputs with other windowing, either {@link Combine.Globally#withoutDefaults

   * withoutDefaults} or {@link Combine.Globally#asSingletonView asSingletonView} must be called.

   *

   * <p>See also {@link #smallest}.

   *

   * <p>See also {@link #of}, which sorts using a user-specified {@code Comparator} function.

   *

   * <p>See also {@link #perKey}, {@link #smallestPerKey}, and {@link #largestPerKey}, which take a

   * {@code PCollection} of {@code KV}s and return the top values associated with each key.

   */

  public static <T extends Comparable<T>> Combine.Globally<T, List<T>> largest(int count) {

    return Combine.globally(largestFn(count));

  }



  /** Returns a {@link TopCombineFn} that aggregates the largest count values. */

  public static <T extends Comparable<T>> TopCombineFn<T, Natural<T>> largestFn(int count) {

    return new TopCombineFn<T, Natural<T>>(count, new Natural<T>()) {};

  }

  /** Returns a {@link TopCombineFn} that aggregates the largest count long values. */

  public static TopCombineFn<Long, Natural<Long>> largestLongsFn(int count) {

    return new TopCombineFn<Long, Natural<Long>>(count, new Natural<Long>()) {};

  }

  /** Returns a {@link TopCombineFn} that aggregates the largest count int values. */

  public static TopCombineFn<Integer, Natural<Integer>> largestIntsFn(int count) {

    return new TopCombineFn<Integer, Natural<Integer>>(count, new Natural<>()) {};

  }

  /** Returns a {@link TopCombineFn} that aggregates the largest count double values. */

  public static TopCombineFn<Double, Natural<Double>> largestDoublesFn(int count) {

    return new TopCombineFn<Double, Natural<Double>>(count, new Natural<>()) {};

  }

  /** Returns a {@link TopCombineFn} that aggregates the smallest count values. */

  public static <T extends Comparable<T>> TopCombineFn<T, Reversed<T>> smallestFn(int count) {

    return new TopCombineFn<T, Reversed<T>>(count, new Reversed<>()) {};

  }

  /** Returns a {@link TopCombineFn} that aggregates the smallest count long values. */

  public static TopCombineFn<Long, Reversed<Long>> smallestLongsFn(int count) {

    return new TopCombineFn<Long, Reversed<Long>>(count, new Reversed<>()) {};

  }

  /** Returns a {@link TopCombineFn} that aggregates the smallest count int values. */

  public static TopCombineFn<Integer, Reversed<Integer>> smallestIntsFn(int count) {

    return new TopCombineFn<Integer, Reversed<Integer>>(count, new Reversed<>()) {};

  }

  /** Returns a {@link TopCombineFn} that aggregates the smallest count double values. */

  public static TopCombineFn<Double, Reversed<Double>> smallestDoublesFn(int count) {

    return new TopCombineFn<Double, Reversed<Double>>(count, new Reversed<>()) {};

  }



  /**

   * Returns a {@code PTransform} that takes an input {@code PCollection<KV<K, V>>} and returns a

   * {@code PCollection<KV<K, List<V>>>} that contains an output element mapping each distinct key

   * in the input {@code PCollection} to the largest {@code count} values associated with that key

   * in the input {@code PCollection<KV<K, V>>}, in decreasing order, sorted using the given {@code

   * Comparator<V>}. The {@code Comparator<V>} must also be {@code Serializable}.

   *

   * <p>If there are fewer than {@code count} values associated with a particular key, then all

   * those values will be in the result mapping for that key, albeit in sorted order.

   *

   * <p>All the values associated with a single key must fit into the memory of a single machine,

   * but there can be many more {@code KV}s in the resulting {@code PCollection} than can fit into

   * the memory of a single machine.

   *

   * <p>Example of use:

   *

   * <pre>{@code

   * PCollection<KV<School, Student>> studentsBySchool = ...;

   * PCollection<KV<School, List<Student>>> top10StudentsBySchool =

   *     studentsBySchool.apply(

   *         Top.perKey(10, new CompareStudentsByAvgGrade()));

   * }</pre>

   *

   * <p>By default, the {@code Coder} of the keys of the output {@code PCollection} is the same as

   * that of the keys of the input {@code PCollection}, and the {@code Coder} of the values of the

   * output {@code PCollection} is a {@code ListCoder} of the {@code Coder} of the values of the

   * input {@code PCollection}.

   *

   * <p>See also {@link #smallestPerKey} and {@link #largestPerKey}, which sort {@code

   * Comparable<V>} values using their natural ordering.

   *

   * <p>See also {@link #of}, {@link #smallest}, and {@link #largest}, which take a {@code

   * PCollection} and return the top elements.

   */

  public static <K, V, ComparatorT extends Comparator<V> & Serializable>

      PTransform<PCollection<KV<K, V>>, PCollection<KV<K, List<V>>>> perKey(

          int count, ComparatorT compareFn) {

    return Combine.perKey(new TopCombineFn<>(count, compareFn));

  }



  /**

   * Returns a {@code PTransform} that takes an input {@code PCollection<KV<K, V>>} and returns a

   * {@code PCollection<KV<K, List<V>>>} that contains an output element mapping each distinct key

   * in the input {@code PCollection} to the smallest {@code count} values associated with that key

   * in the input {@code PCollection<KV<K, V>>}, in increasing order, sorted according to their

   * natural order.

   *

   * <p>If there are fewer than {@code count} values associated with a particular key, then all

   * those values will be in the result mapping for that key, albeit in sorted order.

   *

   * <p>All the values associated with a single key must fit into the memory of a single machine,

   * but there can be many more {@code KV}s in the resulting {@code PCollection} than can fit into

   * the memory of a single machine.

   *

   * <p>Example of use:

   *

   * <pre>{@code

   * PCollection<KV<String, Integer>> keyedValues = ...;

   * PCollection<KV<String, List<Integer>>> smallest10ValuesPerKey =

   *     keyedValues.apply(Top.smallestPerKey(10));

   * }</pre>

   *

   * <p>By default, the {@code Coder} of the keys of the output {@code PCollection} is the same as

   * that of the keys of the input {@code PCollection}, and the {@code Coder} of the values of the

   * output {@code PCollection} is a {@code ListCoder} of the {@code Coder} of the values of the

   * input {@code PCollection}.

   *

   * <p>See also {@link #largestPerKey}.

   *

   * <p>See also {@link #perKey}, which sorts values using a user-specified {@code Comparator}

   * function.

   *

   * <p>See also {@link #of}, {@link #smallest}, and {@link #largest}, which take a {@code

   * PCollection} and return the top elements.

   */

  public static <K, V extends Comparable<V>>

      PTransform<PCollection<KV<K, V>>, PCollection<KV<K, List<V>>>> smallestPerKey(int count) {

    return Combine.perKey(smallestFn(count));

  }



  /**

   * Returns a {@code PTransform} that takes an input {@code PCollection<KV<K, V>>} and returns a

   * {@code PCollection<KV<K, List<V>>>} that contains an output element mapping each distinct key

   * in the input {@code PCollection} to the largest {@code count} values associated with that key

   * in the input {@code PCollection<KV<K, V>>}, in decreasing order, sorted according to their

   * natural order.

   *

   * <p>If there are fewer than {@code count} values associated with a particular key, then all

   * those values will be in the result mapping for that key, albeit in sorted order.

   *

   * <p>All the values associated with a single key must fit into the memory of a single machine,

   * but there can be many more {@code KV}s in the resulting {@code PCollection} than can fit into

   * the memory of a single machine.

   *

   * <p>Example of use:

   *

   * <pre>{@code

   * PCollection<KV<String, Integer>> keyedValues = ...;

   * PCollection<KV<String, List<Integer>>> largest10ValuesPerKey =

   *     keyedValues.apply(Top.largestPerKey(10));

   * }</pre>

   *

   * <p>By default, the {@code Coder} of the keys of the output {@code PCollection} is the same as

   * that of the keys of the input {@code PCollection}, and the {@code Coder} of the values of the

   * output {@code PCollection} is a {@code ListCoder} of the {@code Coder} of the values of the

   * input {@code PCollection}.

   *

   * <p>See also {@link #smallestPerKey}.

   *

   * <p>See also {@link #perKey}, which sorts values using a user-specified {@code Comparator}

   * function.

   *

   * <p>See also {@link #of}, {@link #smallest}, and {@link #largest}, which take a {@code

   * PCollection} and return the top elements.

   */

  public static <K, V extends Comparable<V>> PerKey<K, V, List<V>> largestPerKey(int count) {

    return Combine.perKey(largestFn(count));

  }



  /** @deprecated use {@link Natural} instead */

  @Deprecated

  public static class Largest<T extends Comparable<? super T>>

      implements Comparator<T>, Serializable {

    @Override

    public int compare(T a, T b) {

      return a.compareTo(b);

    }

  }



  /**

   * A {@code Serializable} {@code Comparator} that that uses the compared elements' natural

   * ordering.

   */

  public static class Natural<T extends Comparable<? super T>>

      implements Comparator<T>, Serializable {

    @Override

    public int compare(T a, T b) {

      return a.compareTo(b);

    }

  }



  /** @deprecated use {@link Reversed} instead */

  @Deprecated

  public static class Smallest<T extends Comparable<? super T>>

      implements Comparator<T>, Serializable {

    @Override

    public int compare(T a, T b) {

      return b.compareTo(a);

    }

  }



  /**

   * {@code Serializable} {@code Comparator} that that uses the reverse of the compared elements'

   * natural ordering.

   */

  public static class Reversed<T extends Comparable<? super T>>

      implements Comparator<T>, Serializable {

    @Override

    public int compare(T a, T b) {

      return b.compareTo(a);

    }

  }



  ////////////////////////////////////////////////////////////////////////////



  /**

   * {@code CombineFn} for {@code Top} transforms that combines a bunch of {@code T}s into a single

   * {@code count}-long {@code List<T>}, using {@code compareFn} to choose the largest {@code T}s.

   *

   * @param <T> type of element being compared

   */

  public static class TopCombineFn<T, ComparatorT extends Comparator<T> & Serializable>

      extends AccumulatingCombineFn<T, BoundedHeap<T, ComparatorT>, List<T>>

      implements NameUtils.NameOverride {



    private final int count;

    private final ComparatorT compareFn;



    public TopCombineFn(int count, ComparatorT compareFn) {

      checkArgument(count >= 0, "count must be >= 0 (not %s)", count);

      this.count = count;

      this.compareFn = compareFn;

    }



    @Override

    public String getNameOverride() {

      return String.format("Top(%s)", NameUtils.approximateSimpleName(compareFn));

    }



    @Override

    public BoundedHeap<T, ComparatorT> createAccumulator() {

      return new BoundedHeap<>(count, compareFn, new ArrayList<>());

    }



    @Override

    public Coder<BoundedHeap<T, ComparatorT>> getAccumulatorCoder(

		CoderRegistry registry, Coder<T> inputCoder) {

      return new BoundedHeapCoder<>(count, compareFn, inputCoder);

    }



    @Override

    public void populateDisplayData(DisplayData.Builder builder) {

      super.populateDisplayData(builder);

      builder

          .add(DisplayData.item("count", count).withLabel("Top Count"))

          .add(DisplayData.item("comparer", compareFn.getClass()).withLabel("Record Comparer"));

    }



    @Override

    public String getIncompatibleGlobalWindowErrorMessage() {

      return "Default values are not supported in Top.[of, smallest, largest]() if the input "

          + "PCollection is not windowed by GlobalWindows. Instead, use "

          + "Top.[of, smallest, largest]().withoutDefaults() to output an empty PCollection if the "

          + "input PCollection is empty, or Top.[of, smallest, largest]().asSingletonView() to "

          + "get a PCollection containing the empty list if the input PCollection is empty.";

    }

  }



  /**

   * A heap that stores only a finite number of top elements according to its provided {@code

   * Comparator}. Implemented as an {@link Accumulator} to facilitate implementation of {@link Top}.

   *

   * <p>This class is <i>not</i> safe for multithreaded use, except read-only.

   */

  static class BoundedHeap<T, ComparatorT extends Comparator<T> & Serializable>

      implements Accumulator<T, BoundedHeap<T, ComparatorT>, List<T>> {



    /**

     * A queue with smallest at the head, for quick adds.

     *

     * <p>Only one of asList and asQueue may be non-null.

     */

    @Nullable

	private PriorityQueue<T> asQueue;



    /**

     * A list in with largest first, the form of extractOutput().

     *

     * <p>Only one of asList and asQueue may be non-null.

     */

    @Nullable

	private List<T> asList;



    /** The user-provided Comparator. */

    private final ComparatorT compareFn;



    /** The maximum size of the heap. */

    private final int maximumSize;



    /** Creates a new heap with the provided size, comparator, and initial elements. */

    private BoundedHeap(int maximumSize, ComparatorT compareFn, List<T> asList) {

      this.maximumSize = maximumSize;

      this.asList = asList;

      this.compareFn = compareFn;

    }



    @Override

    public void addInput(T value) {

      maybeAddInput(value);

    }



    /**

     * Adds {@code value} to this heap if it is larger than any of the current elements. Returns

     * {@code true} if {@code value} was added.

     */

    private boolean maybeAddInput(T value) {

      if (maximumSize == 0) {

        // Don't add anything.

        return false;

      }



      // If asQueue == null, then this is the first add after the latest call to the

      // constructor or asList().

      if (asQueue == null) {

        asQueue = new PriorityQueue<>(maximumSize, compareFn);

        for (T item : asList) {

          asQueue.add(item);

        }

        asList = null;

      }



      if (asQueue.size() < maximumSize) {

        asQueue.add(value);

        return true;

      } else if (compareFn.compare(value, asQueue.peek()) > 0) {

        asQueue.poll();

        asQueue.add(value);

        return true;

      } else {

        return false;

      }

    }



    @Override

    public void mergeAccumulator(BoundedHeap<T, ComparatorT> accumulator) {

      for (T value : accumulator.asList()) {

        if (!maybeAddInput(value)) {

          // If this element of accumulator does not make the top N, neither

          // will the rest, which are all smaller.

          break;

        }

      }

    }



    @Override

    public List<T> extractOutput() {

      return asList();

    }



    /** Returns the contents of this Heap as a List sorted largest-to-smallest. */

    private List<T> asList() {

      if (asList == null) {

        List<T> smallestFirstList = Lists.newArrayListWithCapacity(asQueue.size());

        while (!asQueue.isEmpty()) {

          smallestFirstList.add(asQueue.poll());

        }

        asList = Lists.reverse(smallestFirstList);

        asQueue = null;

      }

      return asList;

    }

  }



  /** A {@link Coder} for {@link BoundedHeap}, using Java serialization via {@link CustomCoder}. */

  private static class BoundedHeapCoder<T, ComparatorT extends Comparator<T> & Serializable>

      extends CustomCoder<BoundedHeap<T, ComparatorT>> {

    private final Coder<List<T>> listCoder;

    private final ComparatorT compareFn;

    private final int maximumSize;



    public BoundedHeapCoder(int maximumSize, ComparatorT compareFn, Coder<T> elementCoder) {

      listCoder = ListCoder.of(elementCoder);

      this.compareFn = compareFn;

      this.maximumSize = maximumSize;

    }



    @Override

    public void encode(BoundedHeap<T, ComparatorT> value, OutputStream outStream)

        throws CoderException, IOException {

      listCoder.encode(value.asList(), outStream);

    }



    @Override

    public BoundedHeap<T, ComparatorT> decode(InputStream inStream)

        throws CoderException, IOException {

      return new BoundedHeap<>(maximumSize, compareFn, listCoder.decode(inStream));

    }



    @Override

    public void verifyDeterministic() throws NonDeterministicException {

      verifyDeterministic(this, "HeapCoder requires a deterministic list coder", listCoder);

    }



    @Override

    public boolean isRegisterByteSizeObserverCheap(BoundedHeap<T, ComparatorT> value) {

      return listCoder.isRegisterByteSizeObserverCheap(value.asList());

    }



    @Override

    public void registerByteSizeObserver(

        BoundedHeap<T, ComparatorT> value, ElementByteSizeObserver observer) throws Exception {

      listCoder.registerByteSizeObserver(value.asList(), observer);

    }



    @Override

    public boolean equals(Object other) {

      if (other == this) {

        return true;

      }

      if (!(other instanceof BoundedHeapCoder)) {

        return false;

      }

      BoundedHeapCoder<?, ?> that = (BoundedHeapCoder<?, ?>) other;

      return Objects.equals(this.compareFn, that.compareFn)

          && Objects.equals(this.listCoder, that.listCoder)

          && this.maximumSize == that.maximumSize;

    }



    @Override

    public int hashCode() {

      return Objects.hash(compareFn, listCoder, maximumSize);

    }

  }

}